package com.migration.util;

import com.baomidou.mybatisplus.core.MybatisConfiguration;
import com.baomidou.mybatisplus.core.config.GlobalConfig;
import com.baomidou.mybatisplus.core.metadata.TableInfo;
import com.baomidou.mybatisplus.extension.toolkit.JdbcUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 数据库迁移工具类，提供使用MyBatis-Plus进行数据迁移的辅助功能
 */
@Slf4j
@Component
public class MigrationUtils {

    private final JdbcTemplate oldJdbcTemplate;
    private final JdbcTemplate newJdbcTemplate;
    
    public MigrationUtils(DataSource oldDataSource, DataSource newDataSource) {
        this.oldJdbcTemplate = new JdbcTemplate(oldDataSource);
        this.newJdbcTemplate = new JdbcTemplate(newDataSource);
    }
    
    /**
     * 获取数据库表列表
     *
     * @param dataSource 数据源
     * @param excludePrefixes 要排除的表前缀
     * @return 表名列表
     */
    public List<String> getTableList(DataSource dataSource, List<String> excludePrefixes) {
        List<String> tables = new ArrayList<>();
        
        try (Connection conn = dataSource.getConnection()) {
            DatabaseMetaData metaData = conn.getMetaData();
            String catalog = conn.getCatalog();
            String schema = conn.getSchema();
            
            try (ResultSet rs = metaData.getTables(catalog, schema, "%", new String[]{"TABLE"})) {
                while (rs.next()) {
                    String tableName = rs.getString("TABLE_NAME");
                    
                    // 排除指定前缀的表
                    if (shouldExcludeTable(tableName, excludePrefixes)) {
                        continue;
                    }
                    
                    tables.add(tableName);
                }
            }
        } catch (SQLException e) {
            log.error("获取表列表失败", e);
        }
        
        return tables;
    }
    
    /**
     * 判断表是否应该被排除
     */
    private boolean shouldExcludeTable(String tableName, List<String> excludePrefixes) {
        if (excludePrefixes == null || excludePrefixes.isEmpty()) {
            return false;
        }
        
        for (String prefix : excludePrefixes) {
            if (tableName.startsWith(prefix)) {
                return true;
            }
        }
        
        return false;
    }
    
    /**
     * 获取表的列信息
     *
     * @param dataSource 数据源
     * @param tableName 表名
     * @return 列信息列表
     */
    public List<String> getTableColumns(DataSource dataSource, String tableName) {
        List<String> columns = new ArrayList<>();
        
        try (Connection conn = dataSource.getConnection()) {
            DatabaseMetaData metaData = conn.getMetaData();
            String catalog = conn.getCatalog();
            String schema = conn.getSchema();
            
            try (ResultSet rs = metaData.getColumns(catalog, schema, tableName, "%")) {
                while (rs.next()) {
                    columns.add(rs.getString("COLUMN_NAME"));
                }
            }
        } catch (SQLException e) {
            log.error("获取表 {} 的列信息失败", tableName, e);
        }
        
        return columns;
    }
    
    /**
     * 获取表的总记录数
     *
     * @param jdbcTemplate JdbcTemplate实例
     * @param tableName 表名
     * @return 记录数
     */
    public long getTableRowCount(JdbcTemplate jdbcTemplate, String tableName) {
        String sql = "SELECT COUNT(*) FROM " + tableName;
        return jdbcTemplate.queryForObject(sql, Long.class);
    }
    
    /**
     * 批量迁移表数据
     *
     * @param tableName 表名
     * @param batchSize 每批处理的记录数
     * @return 迁移的记录数
     */
    @Transactional
    public int migrateTableData(String tableName, int batchSize) {
        List<String> columns = getTableColumns(oldJdbcTemplate.getDataSource(), tableName);
        if (columns.isEmpty()) {
            log.error("表 {} 没有列信息", tableName);
            return 0;
        }
        
        String columnsList = String.join(", ", columns);
        String placeholders = columns.stream().map(c -> "?").reduce((a, b) -> a + ", " + b).orElse("");
        
        String selectSql = "SELECT " + columnsList + " FROM " + tableName;
        String insertSql = "INSERT INTO " + tableName + " (" + columnsList + ") VALUES (" + placeholders + ")";
        
        // 由于可能数据量较大，分批读取和写入
        long totalCount = getTableRowCount(oldJdbcTemplate, tableName);
        int processedCount = 0;
        int offset = 0;
        
        while (offset < totalCount) {
            String paginatedSql = selectSql + " LIMIT " + offset + ", " + batchSize;
            List<Map<String, Object>> rows = oldJdbcTemplate.queryForList(paginatedSql);
            
            if (rows.isEmpty()) {
                break;
            }
            
            for (Map<String, Object> row : rows) {
                Object[] params = new Object[columns.size()];
                for (int i = 0; i < columns.size(); i++) {
                    params[i] = row.get(columns.get(i));
                }
                
                newJdbcTemplate.update(insertSql, params);
                processedCount++;
            }
            
            offset += batchSize;
            log.info("表 {} 已迁移 {}/{} 条记录", tableName, processedCount, totalCount);
        }
        
        return processedCount;
    }
    
    /**
     * 迁移单个表
     *
     * @param tableName 表名
     * @return 迁移的记录数
     */
    @Transactional
    public int migrateSingleTable(String tableName) {
        log.info("开始迁移表 {}", tableName);
        try {
            // 使用默认的批处理大小
            return migrateTableData(tableName, 1000);
        } catch (Exception e) {
            log.error("迁移表 {} 失败", tableName, e);
            throw new RuntimeException("迁移表 " + tableName + " 失败: " + e.getMessage(), e);
        }
    }
    
    /**
     * 迁移所有表
     *
     * @param excludePrefixes 要排除的表前缀
     * @return 迁移的表数量
     */
    @Transactional
    public int migrateAllTables(List<String> excludePrefixes) {
        log.info("开始迁移所有表");
        
        List<String> tables = getTableList(oldJdbcTemplate.getDataSource(), excludePrefixes);
        log.info("共发现 {} 个表需要迁移", tables.size());
        
        int migratedTables = 0;
        for (String tableName : tables) {
            try {
                migrateSingleTable(tableName);
                migratedTables++;
            } catch (Exception e) {
                log.error("迁移表 {} 失败，继续迁移其他表", tableName, e);
            }
        }
        
        log.info("成功迁移 {}/{} 个表", migratedTables, tables.size());
        return migratedTables;
    }
    
    /**
     * 获取旧数据库的JdbcTemplate
     */
    public JdbcTemplate getOldJdbcTemplate() {
        return oldJdbcTemplate;
    }
    
    /**
     * 获取新数据库的JdbcTemplate
     */
    public JdbcTemplate getNewJdbcTemplate() {
        return newJdbcTemplate;
    }
} 